﻿using Serilog;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;
using System.Threading.Tasks;
using TorchSharp;
using TorchSharp.Modules;
using static Tensorboard.ApiDef.Types;
using static TorchSharp.torch;
using static TorchSharp.torch.optim;

namespace NgpSharp
{
    public class NeRFTrainer
    {
        public string Name {  get; private set; }
        public Device Device { get; private set; }
        public NeRFNetwork Model { get; private set; }

        public Loss<Tensor, Tensor, Tensor> Criterion { get; private set; }
        public optim.Optimizer Optimizer { get; private set; }
        public lr_scheduler.LRScheduler LRScheduler { get; private set; }
        public ILogger Logger { get; private set; }
        public string Workspace { get; set; }

        private int local_step { get; set; }
        private int global_step { get; set; }
        private int epoch {  get; set; }

        /// <summary>
        /// 
        /// </summary>
        /// <param name="name"></param>
        public NeRFTrainer(string name, Device device, NeRFNetwork model, 
            optim.Optimizer optimizer,
            Loss<Tensor, Tensor, Tensor> criterion, 
            lr_scheduler.LRScheduler scheduler, ILogger logger)
        {
            Name = name;
            Device = device;
            Model = model ?? throw new ArgumentNullException(nameof(model));
            optimizer ??= Adam(model.parameters(), lr: 0.001, weight_decay: 5e-4);
            Optimizer = optimizer;
            Criterion = criterion ?? throw new ArgumentNullException(nameof(criterion));
            scheduler ??= lr_scheduler.LambdaLR(optimizer, (epoch) => 1);
            LRScheduler = scheduler;
            Logger = logger;
        }

        public void Train(NeRFDataLoader loader, int maxEpoch)
        {
            Logger?.Debug("Training begin");
            epoch = 0;
            for (int i = epoch + 1; i < maxEpoch + 1; i++)
            {
                epoch++;
                Logger?.Debug($"begin epoch {i}");
                TrainOneEpoch(loader);
                Logger?.Debug($"epoch {i} finish");
            }
        }

        public void TrainOneEpoch(NeRFDataLoader loader)
        {
            float total_loss = 0;
            float average_loss = 0;
            float loss = 0;
            Model.train();
            local_step = 0;
            var scheduler = LRScheduler;
            var optimizer = Optimizer;

            Logger?.Debug($"start trainning at epoch {epoch}");
            var set = (NeRFDataset)loader.dataset;
            var min = set.BoxMin; var max = set.BoxMax;
            Tensor aabb = torch.from_array(new float[]
            {
                min.X, min.Y, min.Z,
                max.X, max.Y, max.Z,
            });
            foreach (var data in loader)
            {
                local_step += 1;
                global_step += 1;
                optimizer.zero_grad();

                TrainStep(data, aabb, out float outLoss);
                total_loss += outLoss;
                average_loss = total_loss / local_step;
                loss = outLoss;
                optimizer.step();

                float decay_rate = 0.1f;
                float decay_steps = 250 * 1000;
                float new_lrate = (float)(1e-2 * Math.Pow(decay_rate, (global_step / decay_steps)));
                foreach (var item in optimizer.ParamGroups)
                {
                    item.LearningRate = new_lrate;
                }
                Logger?.Debug($"current loss: {loss}; total loss: {total_loss}; average loss: {average_loss}; change lrate: {new_lrate}");
            }
            //scheduler.step();
            Logger?.Debug($"finish epoch {epoch}, final average loss: {average_loss}");
        }
        public void EvaluateOneEpoch(NeRFDataLoader loader)
        {
            Model.eval();
            local_step = 0;

            float total_loss = 0;
            float average_loss = 0;

            Logger?.Debug($"start evaluate at epoch {epoch}");

            var set = (NeRFDataset)loader.dataset;
            var min = set.BoxMin; var max = set.BoxMax;
            Tensor aabb = torch.from_array(new float[]
            {
                min.X, min.Y, min.Z,
                max.X, max.Y, max.Z,
            });
            foreach (var data in loader)
            {
                using (var c = torch.NewDisposeScope())
                {
                    local_step += 1;
                    var result = EvalStep(data, aabb);
                    var loss = result[3];
                    float loss_val = loss.item<float>();
                    total_loss += loss_val;
                }
            }
            average_loss = total_loss / local_step;
            Logger?.Debug($"finish epoch {epoch}, final average loss: {average_loss}");
        }

        /// <summary>
        /// 
        /// </summary>
        /// <param name="data"></param>
        public void TrainStep(Dictionary<string, Tensor> data, Tensor aabb, out float loss_val)
        {
            loss_val = 0;
            using var c = torch.NewDisposeScope();

            var rays_o = data["rays_o"]; //# [B, N, 3]
            var rays_d = data["rays_d"]; //# [B, N, 3]

            var images = data["images"]; //# [B, N, 3/4]

            long B = images.shape[0]; long N = images.shape[1]; long C = images.shape[2];
            Tensor bg_color;
            if (C == 3 || Model.bg_radius > 0) bg_color = 1;
            else bg_color = torch.rand_like(images[TensorIndex.Ellipsis, TensorIndex.Slice(0, 3)], device: rays_o.device);

            Tensor gt_rgb;
            if (C == 4)
            {
                gt_rgb = images[TensorIndex.Ellipsis, TensorIndex.Slice(0, 3)] *
                    images[TensorIndex.Ellipsis, TensorIndex.Slice(3)] +
                    bg_color * (1 - images[TensorIndex.Ellipsis, TensorIndex.Slice(3)]);
            }
            else gt_rgb = images;

            if(aabb.device != rays_o.device) aabb = aabb.to(rays_o.device);
            var outputs = Model.Render(rays_o, rays_d, aabb, staged: false, bg_color: bg_color, perturb: true);

            var pred_rgb = outputs["image"];
            var criterion = Criterion;
            var loss = criterion.call(pred_rgb, gt_rgb).mean([-1]);

            if (loss.shape.Length == 3)
                loss = loss.mean([0]);
            loss = loss.mean();
            loss_val = loss.item<float>();
            loss.backward();
        }
        /// <summary>
        /// 
        /// </summary>
        /// <param name="data"></param>
        /// <returns></returns>
        public List<Tensor> EvalStep(Dictionary<string, Tensor> data, Tensor aabb)
        {
            using var grad = torch.no_grad();
            using var c = torch.NewDisposeScope();

            Tensor rays_o = data["rays_o"]; //# [B, N, 3]
            Tensor rays_d = data["rays_d"]; //# [B, N, 3]
            Tensor images = data["images"]; //# [B, H, W, 3/4]

            long B = images.shape[0];
            long H = images.shape[1];
            long W = images.shape[2];
            long C = images.shape[3];

            Tensor gt_rgb;
            if (C == 4)
            {
                gt_rgb = images[TensorIndex.Ellipsis, TensorIndex.Slice(0, 3)] *
                    images[TensorIndex.Ellipsis, TensorIndex.Slice(3)] +
                    1 * (1 - images[TensorIndex.Ellipsis, TensorIndex.Slice(3)]);
            }
            else gt_rgb = images;

            var outputs = Model.Render(rays_o, rays_d, aabb, staged: true, null, perturb: false);

            Tensor pred_rgb = outputs["image"].reshape(B, H, W, 3);
            Tensor pred_depth = outputs["depth"].reshape(B, H, W);

            var criterion = Criterion;
            Tensor loss = criterion.call(pred_rgb, gt_rgb).mean();

            if (Directory.Exists(Workspace))
            {
                string rootPath = Path.Combine(Workspace, "validation");
                if (!Directory.Exists(rootPath)) Directory.CreateDirectory(rootPath);
                string savePath = Path.Combine(rootPath, $"Epoch{epoch}_{local_step}_rgb.png");
                ExportToImage(pred_rgb, savePath);
                Logger?.Information($"Success export to image with path {savePath}");
            }

            return [pred_rgb.MoveToOuterDisposeScope(), 
                pred_depth.MoveToOuterDisposeScope(), 
                gt_rgb.MoveToOuterDisposeScope(), 
                loss.MoveToOuterDisposeScope()];
        }

        private void ExportToImage(Tensor image, string path)
        {
            if(File.Exists(path)) File.Delete(path);
            image = image.squeeze(0).detach().cpu();
            image = (image * 255).to(torch.ScalarType.Int8);
            using Image<Rgb24> i = new Image<Rgb24>((int)image.shape[1], (int)image.shape[0]);
            i.ProcessPixelRows((acc) =>
            {
                long stri = image.shape[1] * image.shape[2] * image.element_size();
                ReadOnlySpan<byte> bytes = image.bytes;
                for (int i = 0; i < acc.Height; i++)
                {
                    var span = acc.GetRowSpan(i);
                    CopyData(bytes, span, stri, i);
                }
            });
            i.Save(path);
        }
        private void CopyData(ReadOnlySpan<byte> source, Span<Rgb24> target, long length, int step)
        {
            var slice = source.Slice((int)length * step, (int)length);
            ReadOnlySpan<Rgb24> c = MemoryMarshal.Cast<byte, Rgb24>(slice);
            c.CopyTo(target);
        }
    }
}
